import os
import pandas as pd
import shutil

# 文件路径
file_path = "/root/autodl-tmp/FundusReasoner/experiments/med_pub_rl_data/med_pub_rl_train.parquet"
backup_dir = os.path.join(os.path.dirname(file_path), "backup")
backup_path = os.path.join(backup_dir, os.path.basename(file_path))

# 1. 读取 parquet 文件
df = pd.read_parquet(file_path)

# 2. 打印第一个数据的 images 字段
# 假设 images 在每一行是个 list（如果是字符串需用 eval 先转成 list）
print("原 images 字段:", df.iloc[0]['images'])

# 3. 批量替换 images 字段内容
def replace_image_paths(images):
    # images 可能为空或不是 list，要小心处理
    if isinstance(images, list):
        return [img.replace('/home/guoyi/Dataset/public_processed', '/root/autodl-tmp/Dataset/public_processed') for img in images]
    return images

df['images'] = df['images'].apply(replace_image_paths)

# 验证替换
print("替换后 images 字段:", df.iloc[0]['images'])

# 4. 备份原 parquet 文件
os.makedirs(backup_dir, exist_ok=True)
shutil.copy2(file_path, backup_path)
print(f"备份到: {backup_path}")

# 5. 覆盖保存
df.to_parquet(file_path, index=False)
print(f"已覆盖保存到: {file_path}")
